import torch
import torch.nn as nn

from nets.backbone import CBA_M_Block
from nets.backbone import Backbone
from nets.attention import CbamAttention, EcaAttention, SeAttention, CaAttention


attention_block = [SeAttention, CbamAttention, EcaAttention, CaAttention]
# ---------------------------------------------------#
#   yolov3-tiny输出
# ---------------------------------------------------#
def HeadBranch(in_channels, out_channels, detect_channels):
    m = nn.Sequential(
        nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=3, stride=1, padding=1, bias=False),
        nn.BatchNorm2d(num_features=out_channels),
        nn.LeakyReLU(negative_slope=0.1),
        nn.Conv2d(in_channels=out_channels, out_channels=detect_channels, kernel_size=1, stride=1)
    )
    return m

# ---------------------------------------------------#
#   SPP结构，利用不同大小的池化核进行池化
#   池化后堆叠
# ---------------------------------------------------#
class SPP(nn.Module):
    def __init__(self, pool_sizes=[5, 9, 13]):
        super(SPP, self).__init__()
        # 4个池化通道合并
        self.maxpools = nn.ModuleList([nn.MaxPool2d(pool_size, 1, pool_size // 2) for pool_size in pool_sizes])
        # 顺带压缩通道
        self.conv = nn.Conv2d(in_channels=4096, out_channels=1024, kernel_size=1, stride=1, bias=False)

    def forward(self, x):
        features = [maxpool(x) for maxpool in self.maxpools[::-1]]
        features = torch.cat(features + [x], dim=1)
        features = self.conv(features)
        return features


class YoloBody(nn.Module):
    def __init__(self, anchors_mask, num_classes, attention_type, enable_spp, enable_moblile):
        super(YoloBody, self).__init__()

        self.detect_channels = len(anchors_mask[0]) * (5 + num_classes)

        # 注意力机制
        self.attention_type = attention_type
        # 注意力机制
        if self.attention_type >= 1 and self.attention_type <= 4:
            self.feat1_att = attention_block[self.attention_type - 1](1024)
            self.feat2_att = attention_block[self.attention_type - 1](256)
            self.upsample_att = attention_block[self.attention_type - 1](128)

        # SPP开关
        self.enable_spp = enable_spp
        if self.enable_spp:
            self.spp_net = SPP(pool_sizes=[5, 9, 13])

        self.enable_moblile = enable_moblile

        # ---------------------------------------------------#
        #   生成darknet的主干模型
        #   获得两个有效特征层，他们的shape分别是：
        #   26,26,512
        #   13,13,1024
        # ---------------------------------------------------#
        self.backbone = Backbone(self.enable_moblile)


        # enable_maxpool=False  => CBA
        self.cbl2 = CBA_M_Block(in_channels=1024,
                                out_channels=256,
                                kernel_size=1,
                                conv_padding=0,
                                enable_maxpool=False)
        # 转向左预测分支1
        self.head1 = HeadBranch(in_channels=256,
                                out_channels=512,
                                detect_channels=self.detect_channels)

        self.cbl3 = CBA_M_Block(in_channels=256,
                                out_channels=128,
                                kernel_size=1,
                                conv_padding=0,
                                enable_maxpool=False)

        self.up_sample = nn.Upsample(scale_factor=2, mode='nearest')

        # 转向左预测分支2
        self.head2 = HeadBranch(in_channels=384,
                                out_channels=256,
                                detect_channels=self.detect_channels)

    def forward(self, x):
        # ---------------------------------------------------#
        #   获得两个有效特征层，他们的shape分别是：
        #   26,26,256；13,13,512
        # ---------------------------------------------------#
        out1, out2 = self.backbone(x)

        if 1 <= self.attention_type and self.attention_type <= 4:
            out1 = self.feat1_att(out1)
            out2 = self.feat2_att(out2)

        if self.enable_spp:
            out1 = self.spp_net(out1)

        x, _ = self.cbl2(out1)
        predict1 = self.head1(x)

        x, _ = self.cbl3(x)
        x = self.up_sample(x)

        x = torch.cat([out2, x], axis=1)
        predict2 = self.head2(x)

        return predict1, predict2
